### Set working directory to current folder ###
tryCatch({
  # if sourcing in base R
  setwd(getSrcDirectory(function(){})[1])
  # if running in RStudio
  setwd(dirname(rstudioapi::getActiveDocumentContext()$path))
}, error=function(e){})

### Initialize relevant packages ###

library('paramtest')
library('pwr')
library('dplyr')
library('tidyr')
library("ggplot2")

### Estimation Function ###
estBetaParams <- function(mu, sd) {
  var <- sd^2
  alpha <- ((1 - mu) / var - 1 / mu) * mu ^ 2
  beta <- alpha * (1 / mu - 1)
  return(params = c(alpha = alpha, beta = beta))
}

### Test function ###
test_fun <- function(simNum, N, t=12, t_treat = 7, control_mean, control_sd, treat_effect, treat_effect_sd = 0, time_effect, time_effect_sd = 0) {

  require(fixest)
  require(dplyr)

  betaParams <- estBetaParams(control_mean, control_sd)

  control_means <- rbeta(N, betaParams["alpha"], betaParams["beta"])
  time_effects <- rnorm(N, time_effect, time_effect_sd)

  y_control <- sapply(1:t, function(x) {
    mean <- control_means + x*time_effects
    mean[mean < 0] <- 0
    mean[mean > 1] <- 1

    rbinom(N, 1, mean)
  })

  control_means <- rbeta(N, betaParams["alpha"], betaParams["beta"])
  time_effects <- rnorm(N, time_effect, time_effect_sd)
  treat_effects <- rnorm(N, treat_effect, treat_effect_sd)

  y_treated <- sapply(1:t, function(x) {
    mean <- control_means + x*time_effects + (x > t_treat)*treat_effects
    mean[mean < 0] <- 0
    mean[mean > 1] <- 1

    rbinom(N, 1, mean)
  })

  outcome <- c(t(y_control), t(y_treated))
  round <- rep(c(1:t), N*2)
  treatment <- c(rep(0, N * t),
                 rep(c(rep(0, t_treat),
                       rep(1, (t-t_treat))), N))
  id <- rep(1:(N*2), each = t)

  dat <- data.frame(id, round, treatment, outcome)

  # Calculate power for different models
  # model 1
  test_1 <- fixest::feols(outcome ~ treatment,
                          data = dat %>% filter(round %in% c(6,8)),
                          cluster = ~ id)
  beta_1 <- test_1$coeftable["treatment", 1]
  stat_1 <- test_1$coeftable["treatment", 3]
  p_1 <- test_1$coeftable["treatment", 4]

  # model 2
  test_2 <- fixest::feols(outcome ~ treatment | id,
                          data = dat %>% filter(round %in% c(6,8)),
                          cluster = ~ id)
  beta_2 <- test_2$coeftable["treatment", 1]
  stat_2 <- test_2$coeftable["treatment", 3]
  p_2 <- test_2$coeftable["treatment", 4]

  # model 3
  test_3 <- fixest::feols(outcome ~ treatment | id,
                          data = dat,
                          cluster = ~ id)
  beta_3 <- test_3$coeftable["treatment", 1]
  stat_3 <- test_3$coeftable["treatment", 3]
  p_3 <- test_3$coeftable["treatment", 4]

  # model 4
  test_4 <- fixest::feols(outcome ~ treatment + round | id,
                          data = dat,
                          cluster = ~ id)
  beta_4 <- test_4$coeftable["treatment", 1]
  stat_4 <- test_4$coeftable["treatment", 3]
  p_4 <- test_4$coeftable["treatment", 4]

  # model 5
  test_5 <- fixest::feols(outcome ~ treatment | id[round],
                          data = dat,
                          cluster = ~ id)
  beta_5 <- test_5$coeftable["treatment", 1]
  stat_5 <- test_5$coeftable["treatment", 3]
  p_5 <- test_5$coeftable["treatment", 4]

  return (c(beta_1 = beta_1, t_1 = stat_1, p_1 = p_1, sig_1 = (p_1 < .05),
            beta_2 = beta_2, t_2 = stat_2, p_2 = p_2, sig_2 = (p_2 < .05),
            beta_3 = beta_3, t_3 = stat_3, p_3 = p_3, sig_3 = (p_3 < .05),
            beta_4 = beta_4, t_4 = stat_4, p_4 = p_4, sig_4 = (p_4 < .05),
            beta_5 = beta_5, t_5 = stat_5, p_5 = p_5, sig_5 = (p_5 < .05)))
}

## Run power test
power_test <- grid_search(test_fun, n.iter = 500, output = "data.frame",
                          list(N = c(30, 50, 60, 70, 90, 120),
                               control_sd = c(0.2, 0.3, 0.4),
                               treat_effect = c(.05, .1, .15, .2),
                               treat_effect_sd = c(0, 0.02),
                               time_effect = c(-0.01, -0.005, 0, 0.005, 0.01),
                               time_effect_sd = c(0, 0.01)),
                          control_mean = .7,
                          t = 12, t_treat = 7,
                          parallel = "multicore", ncpus = 20)

## Clean up results
power <- results(power_test) %>%
  group_by(N.test, control_sd.test,
           treat_effect.test, treat_effect_sd.test,
           time_effect.test, time_effect_sd.test) %>%
  mutate(across(where(is.character), as.numeric)) %>%
  summarise(power_1 = mean(sig_1, na.rm = TRUE),
            power_2 = mean(sig_2, na.rm = TRUE),
            power_3 = mean(sig_3, na.rm = TRUE),
            power_4 = mean(sig_4, na.rm = TRUE),
            power_5 = mean(sig_5, na.rm = TRUE)) %>%
  ungroup %>%
  pivot_longer(cols = starts_with("power"),
               names_to = "model",
               names_prefix = "power_",
               names_transform = list (model = as.integer),
               values_to = "power")

## Figure C8 ##
## Plot showing effect on power of treatment effect size across all five models
## Assuming no trends
ggplot(power %>% filter(time_effect.test == 0,
                        treat_effect_sd.test == 0.02,
                        time_effect_sd.test == 0),
       aes(x = N.test,
           y = power,
           group = factor(treat_effect.test),
           colour = factor(treat_effect.test))) +
  geom_point(size = 1) +
  geom_line() +
  geom_hline(yintercept = 0.8, linetype = "dotted") +
  geom_vline(xintercept = 30, linetype = "dotted") +
  labs(x='Sample Size (N per group)', y='Power (%)', colour="Treatment Effect") +
  scale_y_continuous(labels = scales::percent_format(suffix = ""),
                     breaks = seq(0, 1, .2),
                     limits = c(0,1),
                     sec.axis = sec_axis(~ . ,
                                         name = "Respondent Heterogeneity\n(Indvidual Mean SD)",
                                         breaks = NULL,
                                         labels = NULL)) +
  scale_x_continuous(breaks = seq(0, 120, 30),
                     sec.axis = sec_axis(~ . ,
                                         name = "Model",
                                         breaks = NULL,
                                         labels = NULL)) +
  scale_color_manual(values = c("#cccccc",
                                "#969696",
                                "#636363",
                                "#252525")) +
  facet_grid(control_sd.test ~ model) +
  theme_minimal() +
  theme(legend.position = "bottom",
        strip.text = element_text(size = 10),
        axis.text.x = element_text(size = 8),
        axis.text.y = element_text(size = 8))
ggsave("fig_c8.png",
       width = 7, height = 6, units = "in",
       bg = "white")

## Figure C9 ##
## Plot showing effect on power of treatment effect size across all five models
## Assuming time trends
ggplot(power %>% filter(time_effect.test == -0.005,
                        treat_effect_sd.test == 0.02,
                        time_effect_sd.test == 0),
       aes(x = N.test,
           y = power,
           group = factor(treat_effect.test),
           colour = factor(treat_effect.test))) +
  geom_point(size = 1) +
  geom_line() +
  geom_hline(yintercept = 0.8, linetype = "dotted") +
  geom_vline(xintercept = 30, linetype = "dotted") +
  labs(x='Sample Size (N per group)', y='Power (%)', colour="Treatment Effect") +
  scale_y_continuous(labels = scales::percent_format(suffix = ""),
                     breaks = seq(0, 1, .2),
                     limits = c(0,1),
                     sec.axis = sec_axis(~ . ,
                                         name = "Respondent Heterogeneity\n(Indvidual Mean SD)",
                                         breaks = NULL,
                                         labels = NULL)) +
  scale_x_continuous(breaks = seq(0, 120, 30),
                     sec.axis = sec_axis(~ . ,
                                         name = "Model",
                                         breaks = NULL,
                                         labels = NULL)) +
  scale_color_manual(values = c("#cccccc",
                                "#969696",
                                "#636363",
                                "#252525")) +
  facet_grid(control_sd.test ~ model) +
  theme_minimal() +
  theme(legend.position = "bottom",
        strip.text = element_text(size = 10),
        axis.text.x = element_text(size = 8),
        axis.text.y = element_text(size = 8))
ggsave("fig_c9.png",
       width = 7, height = 6, units = "in",
       bg = "white")

